import os
from typing import Any
import httpx
from mcp.server.fastmcp import FastMCP
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.client.sse import sse_client
import json
import datetime
from openai import AsyncOpenAI
from dotenv import load_dotenv

# Initialize FastMCP server
mcp = FastMCP("weather")

AMAP_KEY = os.getenv("AMAP_KEY", "3bc7e2adbbf9686cab5a4daf3484b498")
AMAP_API_BASE = "https://restapi.amap.com/v3"
AMAP_API_V5_BASE = "https://restapi.amap.com/v5"

async def amap_request(endpoint: str, params: dict) -> dict[str, Any] | None:
    params["key"] = AMAP_KEY
    url = f"{AMAP_API_BASE}/{endpoint}"
    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(url, params=params, timeout=20.0)
            response.raise_for_status()
            return response.json()
        except Exception:
            return None

async def amap_v5_request(endpoint: str, params: dict) -> dict[str, Any] | None:
    params["key"] = AMAP_KEY
    url = f"{AMAP_API_V5_BASE}/{endpoint}"
    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(url, params=params, timeout=20.0)
            response.raise_for_status()
            return response.json()
        except Exception:
            return None

@mcp.tool()
async def amap_geo(address: str, city: str = None) -> str:
    """
    地理编码: 地址转经纬度
    Args:
        address (str): 需要查询的地址。
        city (str, 可选): 城市名，用于限定地址解析范围。
    Returns:
        str: 经纬度字符串，格式为 '经度,纬度'
    """
    params = {"address": address}
    if city:
        params["city"] = city
    data = await amap_request("geocode/geo", params)
    if not data or data.get("status") != "1" or not data.get("geocodes"):
        return "地理编码失败"
    # 只返回location字段
    return data["geocodes"][0].get("location", "地理编码失败")

@mcp.tool()
async def amap_regeocode(location: str) -> str:
    """
    逆地理编码: 经纬度转地址
    Args:
        location (str): 经纬度，格式为 '经度,纬度'。
    Returns:
        str: 格式化后的地址
    """
    params = {"location": location}
    data = await amap_request("geocode/regeo", params)
    if not data or data.get("status") != "1" or not data.get("regeocode"):
        return "逆地理编码失败"
    # 只返回formatted_address字段
    return data["regeocode"].get("formatted_address", "逆地理编码失败")

@mcp.tool()
async def amap_direction_driving(origin: str, destination: str) -> str:
    """
    驾车路径规划
    Args:
        origin (str): 起点经纬度，格式为 '经度,纬度'。
        destination (str): 终点经纬度，格式为 '经度,纬度'。
    Returns:
        str: 完整的API响应数据
    """
    params = {"origin": origin, "destination": destination}
    data = await amap_request("direction/driving", params)
    if not data or data.get("status") != "1":
        return "驾车路径规划失败"
    return json.dumps(data, ensure_ascii=False, indent=2)



@mcp.tool()
async def amap_weather(city: str, extensions: str = "all") -> str:
    """
    天气查询
    Args:
        city (str): 城市名称。
        extensions (str, 可选): 气象类型，可选值：base/all，默认为base。
            - base: 返回实况天气
            - all: 返回预报天气
    Returns:
        str: 完整的API响应数据
    """
    params = {"city": city, "extensions": extensions}
    data = await amap_request("weather/weatherInfo", params)
    if not data or data.get("status") != "1":
        return "天气查询失败"
    return json.dumps(data, ensure_ascii=False, indent=2)

@mcp.tool()
async def amap_around_search(keywords: str, location: str, radius: str = "1000") -> str:
    """
    周边搜索
    Args:
        keywords (str): 搜索关键词，如"美食"、"酒店"等。
        location (str): 中心点经纬度，格式为 '经度,纬度'。
        radius (str, 可选): 搜索半径，单位为米，默认1000。
    Returns:
        str: 完整的API响应数据
    """
    params = {"keywords": keywords, "location": location, "radius": radius}
    data = await amap_request("place/around", params)
    if not data or data.get("status") != "1":
        return "周边搜索失败"
    return json.dumps(data, ensure_ascii=False, indent=2)

@mcp.tool()
async def amap_distance(origins: str, destination: str, type: str = "1") -> str:
    """
    距离测量
    Args:
        origins (str): 起点经纬度，格式为 '经度,纬度'。
        destination (str): 终点经纬度，格式为 '经度,纬度'。
        type (str, 可选): 距离类型，1为驾车，0为直线，3为步行，默认1。
    Returns:
        str: 完整的API响应数据
    """
    params = {"origins": origins, "destination": destination, "type": type}
    data = await amap_request("distance", params)
    if not data or data.get("status") != "1":
        return "距离测量失败"
    return json.dumps(data, ensure_ascii=False, indent=2)

# 添加ip定位
@mcp.tool()
async def get_ip_location() -> str:
    """
    获取IP地址的地理位置信息
    Returns:
        str: 格式化后的位置信息，包含城市和详细位置
    """
    params = {}  # 不需要额外参数，默认使用请求的IP
    data = await amap_request("ip", params)
    if not data or data.get("status") != "1":
        return "IP定位失败，将使用默认地址：东莞市"
    
    province = data.get("province", "")
    city = data.get("city", "")
    district = data.get("district", "")
    adcode = data.get("adcode", "")
    rectangle = data.get("rectangle", "")  # 经纬度范围
    
    # 提取中心点经纬度（取rectangle的中间值）
    location = ""
    if rectangle and "," in rectangle:
        try:
            # rectangle格式是"左下经度,左下纬度;右上经度,右上纬度"
            parts = rectangle.split(';')
            if len(parts) == 2:
                ll = [float(x) for x in parts[0].split(',')]  # 左下角
                ur = [float(x) for x in parts[1].split(',')]  # 右上角
                # 计算中心点
                center_lng = (ll[0] + ur[0]) / 2
                center_lat = (ll[1] + ur[1]) / 2
                location = f"{center_lng:.6f},{center_lat:.6f}"
        except:
            pass
    
    result = {
        "status": "1",
        "info": "IP定位成功",
        "province": province,
        "city": city,
        "district": district,
        "adcode": adcode,
        "location": location
    }
    
    return json.dumps(result, ensure_ascii=False, indent=2)

@mcp.tool()
async def get_current_time(time: str = "default") -> str:
    """
    获取当前计算机的时间
    Args:
        time (str, 可选): 时间参数，默认值为"default"。
    Returns:
        str: 当前时间，格式为 'YYYY-MM-DD HH:MM:SS'
    """
    current_time = datetime.datetime.now()
    return current_time.strftime("%Y-%m-%d %H:%M:%S")

@mcp.tool()
async def amap_text_search(keywords: str, city: str = None, page_num: str = "1", page_size: str = "10") -> str:
    """
    POI关键词搜索
    Args:
        keywords (str): 搜索关键词，如"美食"、"酒店"、"景点"等。
        city (str, 可选): 指定城市名，如"上海"、"北京"等。
        page_num (str, 可选): 当前页数，默认1。
        page_size (str, 可选): 每页记录数，默认10，最大20。
    Returns:
        str: 完整的API响应数据，包含POI名称、地址、坐标、图片等信息
    """
    params = {
        "keywords": keywords,
        "page_num": page_num,
        "page_size": page_size,
        "show_fields": "business,photos,deep_info,richtext"  # 返回丰富字段
    }
    if city:
        params["region"] = city
    
    data = await amap_v5_request("place/text", params)
    if not data or data.get("status") != "1":
        return f"POI关键词搜索失败，无法查找'{keywords}'"
    
    # 简化输出，保留最重要的字段
    formatted_data = {
        "status": "1",
        "count": data.get("count", "0"),
        "pois": []
    }
    
    for poi in data.get("pois", []):
        formatted_poi = {
            "id": poi.get("id", ""),  # 添加ID字段以便后续查询详情
            "name": poi.get("name", ""),
            "address": poi.get("address", ""),
            "location": poi.get("location", ""),
            "type": poi.get("type", ""),
            "tel": poi.get("tel", ""),
            "rating": poi.get("rating", ""),
            "photos": [p.get("url") for p in poi.get("photos", [])[:3]] if "photos" in poi else []
        }
        formatted_data["pois"].append(formatted_poi)
    
    return json.dumps(formatted_data, ensure_ascii=False, indent=2)

@mcp.tool()
async def amap_place_detail(id: str) -> str:
    """
    POI详细信息查询
    Args:
        id (str): POI ID，通常从POI关键词搜索结果中获取。
    Returns:
        str: POI的详细信息，包含详细地址、联系方式、营业时间、评分、图片等
    """
    params = {
        "id": id,
        "show_fields": "business,photos,deep_info,richtext,indoor_data"  # 返回所有可用字段
    }
    
    data = await amap_v5_request("place/detail", params)
    if not data or data.get("status") != "1" or not data.get("pois"):
        return f"POI详情查询失败，ID: {id}"
    
    poi = data.get("pois", [])[0]
    
    # 构建更友好的返回格式
    formatted_detail = {
        "id": poi.get("id", ""),
        "name": poi.get("name", ""),
        "address": poi.get("address", ""),
        "location": poi.get("location", ""),
        "tel": poi.get("tel", ""),
        "website": poi.get("website", ""),
        "email": poi.get("email", ""),
        "rating": poi.get("rating", ""),
        "tag": poi.get("tag", ""),
        "type": poi.get("type", ""),
        "business_area": poi.get("business_area", ""),
        "open_time": poi.get("open_time", ""),
        "photos": []
    }
    
    # 提取照片
    if "photos" in poi and poi["photos"]:
        for photo in poi["photos"][:5]:  # 最多取5张照片
            if "url" in photo:
                formatted_detail["photos"].append(photo["url"])
    
    # 提取丰富文本描述
    if "richtext" in poi and poi["richtext"]:
        if "detail" in poi["richtext"]:
            formatted_detail["description"] = poi["richtext"]["detail"]
            
        if "transit" in poi["richtext"]:
            formatted_detail["transit_info"] = poi["richtext"]["transit"]
    
    # 提取深度信息
    if "deep_info" in poi and poi["deep_info"]:
        if "ticket_info" in poi["deep_info"]:
            formatted_detail["ticket_info"] = poi["deep_info"]["ticket_info"]
            
        if "opening_hours" in poi["deep_info"]:
            formatted_detail["opening_hours"] = poi["deep_info"]["opening_hours"]
    
    return json.dumps(formatted_detail, ensure_ascii=False, indent=2)


@mcp.tool()
def summarize_route_plan(subtasks: str) -> str:
    """
    汇总所有子任务及工具结果，生成最终的路线规划文本。
    Args:
        subtasks (str): 包含原始问题，子问题，子问题答案等。
    Returns:
        str: 路线规划文本
    """
    import os
    import openai
    load_dotenv()
    openai.api_key = os.environ.get('OPENAI_API_KEY', '356a47d7-0202-4c9d-87c4-547cfc1a1a95')
    base_url = os.environ.get('OPENAI_BASE_URL', None)
    if base_url:
        openai.base_url = base_url
    model = os.environ.get('OPENAI_MODEL', 'Qwen/Qwen3-235B-A22B')
    prompt = "请根据以下子任务的工具结果和总结，生成最终的规划安排，必须包含图片的url，不得删去，方便可视化观看：\n"
    prompt += subtasks
    messages = [
        {"role": "system", "content": "你是一个帮助用户整合多步骤任务结果并给出最终规划的智能助手，主要通过时间流程去安排日程。"},
        {"role": "user", "content": prompt}
    ]
    response = openai.chat.completions.create(
        model=model,
        messages=messages
    )
    return response.choices[0].message.content

# @mcp.tool()
# def extract_route_points(final_plan: str) -> str:
#     """
#     从路线规划文本中提取行程点列表。
#     Args:
#         final_plan (str): 路线规划文本
#     Returns:
#         str: 返回地址组成的字符串
#     """
#     import os
#     import openai
#     load_dotenv()
#     openai.api_key = os.environ.get('OPENAI_API_KEY', '356a47d7-0202-4c9d-87c4-547cfc1a1a95')
#     base_url = os.environ.get('OPENAI_BASE_URL', 'https://api-inference.modelscope.cn/v1/')
#     if base_url:
#         openai.base_url = base_url
#     model = os.environ.get('OPENAI_MODEL', 'Qwen/Qwen3-235B-A22B')
#     extract_points_prompt = (
#         "请从以下行程规划内容中，按照时间顺序提取出所有旅游路线点，以list形式返回，每个点包含停留的地点名称，顺序为行程顺序。\n"
#         "【示例】\n"
#         "用户行程：我从东莞出发，中午在松山湖吃饭，下午到深圳。\n"
#         "输出：\n"
#         "\n 东莞,松山湖,深圳\n"
#         "\n请严格按照上述格式提取：必须是实际的地点名称，包括城市名称，景点名称，酒店地址，餐厅地址\n" + final_plan
#     )
#     extract_messages = [
#         {"role": "system", "content": "你是一个位置信息提取助手。"},
#         {"role": "user", "content": extract_points_prompt}
#     ]
#     print("extract_messages内容：", extract_messages)
#
#     extract_response = openai.chat.completions.create(
#         model=model,
#         messages=extract_messages
#     )
#     print("extract_response内容：", extract_response)
#     return extract_response.choices[0].message.content
#
class MCPClient:
    def __init__(self, config_path):
        self.sessions = {}  # {server_key: session}
        self.exit_stack = AsyncExitStack()
        self.config = self.load_config(config_path)

    def load_config(self, path):
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)

    async def connect_all(self):
        for key, conf in self.config.items():
            if conf["type"] == "local":
                server_params = StdioServerParameters(
                    command=conf["command"],
                    args=conf["args"],
                    env=None
                )
                stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
                session = await self.exit_stack.enter_async_context(ClientSession(*stdio_transport))
                await session.initialize()
                self.sessions[key] = session
            elif conf["type"] == "sse":
                sse_transport = await self.exit_stack.enter_async_context(sse_client(conf["url"]))
                session = await self.exit_stack.enter_async_context(ClientSession(*sse_transport))
                await session.initialize()
                self.sessions[key] = session

    async def list_tools(self):
        tools = []
        for key, session in self.sessions.items():
            for tool in (await session.list_tools()).tools:
                tool.name = f"{key}_{tool.name}"
                tools.append(tool)
        return tools

    async def call_tool(self, tool_name, args):
        for key, session in self.sessions.items():
            prefix = f"{key}_"
            if tool_name.startswith(prefix):
                return await session.call_tool(tool_name[len(prefix):], args)
        raise Exception("Unknown tool prefix")

if __name__ == "__main__":
    # Initialize and run the server
    mcp.run(transport='stdio')